"""
Copyright © 2024 Walkline Wang (https://walkline.wang)
Gitee: https://gitee.com/walkline/grouping-teams
"""
import sys, os
import csv, json
import math
import random


__VERSION__ = 'v0.0.3'

TEAMS_PER_GROUP = 20
'''分组原则：每组最多20支队伍，尽量不出现同一所学校的两支队伍在同一组的情况'''
MAX_RETRY_COUNTS = 100
'''根据分组原则此题没有最优解，因此随机分组最多尝试 MAX_RETRY_COUNTS 次'''

METHOD_1 = 0
'''复赛、省赛模式，尽量避免同校竞争，分组数量会增加'''
METHOD_2 = 1
'''校赛模式，随机抽取队伍分组，分组数量最少'''
METHOD_3 = 2
'''生成或更新 JSON 数据文件'''


class Team(object):
	def __init__(self, index: str, belongs_to, name, leader, member_1, member_2, leader_account):
		self.index = index
		self.belongs_to = belongs_to
		self.name = name
		self.leader = leader
		self.member_1 = member_1
		self.member_2 = member_2
		self.leader_account = leader_account

	def __getitem__(self) -> list:
		return vars(self).values()


#region functions
def show_help():
	usage =\
f'''
{'=' * (24 + len(__VERSION__) + 1 if len(__VERSION__) else 0)}
  \033[1;37m参赛队伍自动分组工具\033[0m {__VERSION__}
{'=' * (24 + len(__VERSION__) + 1 if len(__VERSION__) else 0)}
使用方法：
    python3 main.py [CSV_FILE]

    CSV 文件格式：
        团队序号, 学校名称, 团队名称, 队长姓名, 队员1姓名, 队员2姓名, 队长账号
'''

	print(usage)

def choose_an_option(title: str, options: list):
	print(f'\n{title} List:')
	for index, option in enumerate(options, start=1):
		if index == 1:
			print(f'\x1b[32m  [{index}] {option}\033[0m')
		else:
			print(f'  [{index}] {option}')

	selected = None

	while True:
		try:
			selected = input('Choose an option: ')

			if selected == '':
				return 0

			selected = int(selected)

			assert type(selected) is int and 0 < selected <= len(options)
			
			return selected - 1
		except KeyboardInterrupt:
			exit()
		except:
			pass

def read_and_analyse_data(filename, method: int = METHOD_1) -> tuple:
	'''
	读取 csv 文件中的数据，并统计所有学校中最多队伍的数量。

	@param filename: 要读取的文件名。
	@param method: 处理数据的方法，默认为 METHOD_1。
	@return: 包含两个元素的元组，第一个元素是队伍列表，第二个元素是最多队伍数量的学校数量。
	'''
	teams = []
	schools = []

	with open(filename, 'r') as csv_file:
		csv_reader = csv.reader(csv_file)

		for row in csv_reader:
			if row[0].isdigit():
				teams.append(Team(row[0], row[1], row[2], row[3], row[4], row[5], row[6]) if method == METHOD_1 else row)
				schools.append(row[1])

	return teams, max(dict([[i, schools.count(i)] for i in set(schools)]).values())

def convert_csv_to_json(filename) -> tuple:
	'''
	将分组后输出的 CSV 文件转换为 JSON 文件。
	分组文件可能需要手动调整数据，可以使用此函数重新生成 JSON 数据文件。

	@param filename: 要转换的文件名。
	'''
	groups = []
	remaining_teams = []
	is_remaining = False

	with open(filename, 'r') as csv_file:
		csv_reader = csv.reader(csv_file)
		current_teams = []

		for row in csv_reader:
			if row[0].startswith("Group ") or row[0].startswith("Remaining"):
				is_remaining = True if row[0].startswith("Remaining") else False
				current_teams = []
			elif row[0] == '':
				if is_remaining:
					remaining_teams.extend(current_teams)
				else:
					groups.append(current_teams)
			else:
				current_teams.append(Team(row[0], row[1], row[2], row[3], row[4], row[5], row[6]))

	if groups[-1] != current_teams:
		groups.append(current_teams)

	return groups, remaining_teams

def calculate_teams_limit(teams: list, least_groups_count: int, method: int = METHOD_1) -> int:
	'''
	返回相对合理的分组队伍数量。

	@param teams: 所有队伍的列表。
	@param least_groups_count: 最少分组数。
	@param method: 分组方式，默认为 METHOD_1。
	@return: 推荐的每组队伍数量。
	'''
	limits = []
	results = []

	teams_count = len(teams)
	groups_count = math.ceil(teams_count / TEAMS_PER_GROUP)
	total = max(groups_count, least_groups_count) if method == METHOD_1 else groups_count

	for limit in range(1, TEAMS_PER_GROUP + 1):
		result = teams_count / limit

		if result <= total:
			limits.append(limit)
			results.append(result)

	return limits[results.index(max(results))]

def output_group_data(separated_groups: list, remaining_teams: list, filename, method: int = METHOD_1):
	'''输出分组数据，CSV 文件用于查看分组详情，JSON 文件用于自动化处理。'''
	letter_table = [chr(i) for i in range(65, 91)]
	filename, ext = os.path.splitext(filename)
	json_dict = {}

	with open(f'{filename}-output{ext}', 'w', newline='') as csv_file:
		csv_writer = csv.writer(csv_file)

		if remaining_teams:
			teams_list = []
			csv_writer.writerow(['Remaining', f'{len(remaining_teams)} teams'])

			for team in remaining_teams:
				teams_list.append(team.leader_account if method == METHOD_1 else team[6])
				csv_writer.writerow(team.__getitem__() if method == METHOD_1 else team)

			csv_writer.writerow(['', ''])
			json_dict['Remainings'] = teams_list

		for index, teams in enumerate(separated_groups):
			teams_list = []
			csv_writer.writerow([f'Group {letter_table[index]}', f'{len(teams)} teams'])

			for team in teams:
				teams_list.append(team.leader_account if method == METHOD_1 else team[6])
				csv_writer.writerow(team.__getitem__() if method == METHOD_1 else team)

			csv_writer.writerow(['', ''])
			json_dict[f'Group_{letter_table[index]}___REPLACE_WITH_GAME_ID'] = teams_list

	with open(f'{filename}-output.json', 'w') as json_file:
		json.dump(json_dict, json_file)

def show_summary(teams: list, separated_groups: list, remaining_teams: list):
	'''汇总分组情况。'''
	teams_count = 0
	letter_table = [chr(i) for i in range(65, 91)]

	print(f'Teams: {len(teams)}\nGroups: {len(separated_groups)}')

	for index, teams in enumerate(separated_groups):
		teams_count += len(teams)
		print(f'  - Group {letter_table[index]} has {len(teams)} teams')
	print(f'  - Total: {teams_count} teams')

	print(f'Remaining {len(remaining_teams)} team(s)')

def separating_groups(raw_teams: list, groups_count: int, teams_limit: int = TEAMS_PER_GROUP, method: int = METHOD_1) -> tuple:
	'''
	将原始队伍列表分成指定数量的小组，每个小组最多包含指定数量的队伍。

	:param raw_teams: 原始队伍列表。
	:param groups_count: 需要分组的数量。
	:param teams_limit: 每个小组的最大队伍数量，默认为 TEAMS_PER_GROUP。
	:param method: 分组方式，默认为 METHOD_1。
	:return: 返回已分组列表和未分组列表。
	'''
	random.shuffle(raw_teams)

	remaining_teams = list(raw_teams)
	groups = [] # 分组结果

	for _ in range(groups_count):
		current_teams = []
		current_schools = []
		current_teams_name = []
		retry_count = 0

		while len(current_teams) < teams_limit and remaining_teams:
			lucky_team = random.choice(list(remaining_teams))

			if method == METHOD_1:
				if lucky_team.belongs_to not in current_schools and lucky_team.name not in current_teams_name:
					current_schools.append(lucky_team.belongs_to)
					current_teams_name.append(lucky_team.name)
					current_teams.append(lucky_team)
					remaining_teams.remove(lucky_team)
				else:
					if retry_count > MAX_RETRY_COUNTS:
						groups.append(current_teams)
						break

					retry_count += 1
			else:
				current_teams.append(lucky_team)
				remaining_teams.remove(lucky_team)

		if retry_count <= MAX_RETRY_COUNTS:
			groups.append(current_teams)

	return groups, list(remaining_teams)
#endregion


def main(csv_file: str):
	method = choose_an_option('选择分组方式', ['复赛、省赛（避免同校竞争，分组数量较多）', '校赛（分组数量最少）', '生成/更新 JSON 数据文件'])

	if method == METHOD_3:
		separated_groups, remaining_teams = convert_csv_to_json(csv_file)
		output_group_data(separated_groups, remaining_teams, csv_file)
		return

	teams, least_groups_count = read_and_analyse_data(csv_file, method)
	teams_limit = calculate_teams_limit(teams, least_groups_count, method)
	groups_count = math.ceil(len(teams) / teams_limit)

	retry_count = MAX_RETRY_COUNTS

	while retry_count > 0:
		separated_groups, remaining_teams = separating_groups(teams, groups_count, teams_limit, method)

		if len(remaining_teams) == 0:
			break

		retry_count -= 1

	output_group_data(separated_groups, remaining_teams, csv_file, method)
	show_summary(teams, separated_groups, remaining_teams)


if __name__ == '__main__':
	if len(sys.argv) > 1:
		csv_file = sys.argv[1]

		if os.path.exists(csv_file):
			main(csv_file)
		else:
			print(f'File {csv_file} not exists!')
	else:
		show_help()
